import os
import json
from openai import OpenAI
from dotenv import load_dotenv


load_dotenv()  # 加载环境变量


base_url = "https://dashscope.aliyuncs.com/compatible-mode/v1"
model = "qwen3-32b"
client = OpenAI(
    api_key=os.getenv("BAILIAN_KEY"),
    base_url=base_url,
)


def add_numbers(a: int, b: int) -> int:
    return a + b


# 定义工具（函数）描述
tools = [
    {
        "type": "function",
        "function": {
            "name": "add_numbers",
            "description": "计算两个整数的和",
            "parameters": {
                "type": "object",
                "properties": {
                    "a": {"type": "integer", "description": "第一个整数"},
                    "b": {"type": "integer", "description": "第二个整数"},
                },
                "required": ["a", "b"],
            },
        },
    }
]


def process_query(query: str) -> str:
    messages = [{"role": "user", "content": query}]

    print(f"\n[1]REQ: {messages}\n")
    responses = client.chat.completions.create(
        model=model,
        messages=messages,
        tools=tools,
        tool_choice="auto",  # 让模型自动决定是否调用函数
    )
    response = responses.choices[0]
    print(f"[1]RES: {response.model_dump()}\n")

    if response.finish_reason == "tool_calls":
        messages.append(response.message.model_dump())
        for tool_call in response.message.tool_calls:
            func_name = tool_call.function.name
            func_args = json.loads(tool_call.function.arguments)

            if func_name == "add_numbers":
                result = add_numbers(**func_args)
                messages.append(
                    {
                        "role": "tool",
                        "content": f"{result}",
                        "tool_call_id": tool_call.id,
                    }
                )

        print(f"[2]REQ: {messages}\n")
        responses = client.chat.completions.create(model=model, messages=messages)
        response = responses.choices[0]
        print(f"[2]RES: {response.model_dump()}\n")

        return response.message.content

    elif response.finish_reason == "stop":
        return response.message.content


query = "13与29之和"
print("问:", query)
answer = process_query(query)
print("答:", answer)
